

import pennylane as qml
from pennylane.devices import Device, DefaultExecutionConfig
from pennylane.tape import QuantumScript, QuantumScriptOrBatch
from pennylane.devices.execution_config import DefaultExecutionConfig, ExecutionConfig
from pennylane.transforms.core import TransformProgram
from dataclasses import replace
import sys
#from pyqos.experiment.data_taking.scan_circuits import RunCircuits
import numpy as np
import math
class Quanta_CS(Device):
    """
        
    """
    @property
    def name(self):
        return "QuantaCS.superc"
    
    @property
    def wires(self):
        return qml.wires.Wires(self.qubits)
    
    def __init__(self, wires, datatype="P01", printOP = False):
        super(Quanta_CS, self).__init__(wires=wires)
        self.qubits = wires
        self.coupler = []
        self.readline = []
        self.datatype = datatype
        self.instruction=[]#生成QCIS指令
        self.queue = []      # 存储操作队列
        self.measurements = []  # 存储测量指令
        self.printOP = printOP
        # self.compiler = Quanta_CS(qubits=self.qubits)
        

    
    def check_qubit(self,qubit):#检验qubit是否在范围内
        if qubit<0 or qubit>65:
            raise ValueError(f"The qubit{qubit} is out of range, please select the qubit again.")
            sys.exit()
        elif qubit%1!=0:
            raise ValueError(f"The qubit{qubit} is not an integer, please select the qubit again.")
    
    def chip_topology(self,qubit1, qubit2):#对两比特门是否能耦合
        if abs(qubit1-qubit2) not in [5,6,7]:
            raise ValueError(f"There is no coupler between qubit{qubit1} and qubit{qubit2}, please select the qubits again.")
            sys.exit()
        elif qubit1 in [0,1,2,3,4] or qubit2 in [0,1,2,3,4]:
            if abs(qubit1-qubit2) not in [6,7]:
                raise ValueError(f"There is no coupler between qubit{qubit1} and qubit{qubit2}, please select the qubits again.")
                sys.exit()
        elif qubit1 in [5,6,17,18,29,30,41,42,53,54,65] or qubit2 in [5,6,17,18,29,30,41,42,53,54,65]:
            if abs(qubit1-qubit2) not in [6]:
                raise ValueError(f"There is no coupler between qubit{qubit1} and qubit{qubit2}, please select the qubits again.")
                sys.exit()
        elif qubit1 in [60,61,62,63,64] or qubit2 in [60,61,62,63,64]:
            if abs(qubit1-qubit2) not in [5,6]:
                raise ValueError(f"There is no coupler between qubit{qubit1} and qubit{qubit2}, please select the qubits again.")
                sys.exit()
    
    def coupler_(self,qubit1, qubit2):#B:G
        if qubit2>=qubit1:
            qubit2,qubit1=qubit1,qubit2
        c_now='G'+str(qubit1).zfill(2)+str(qubit2).zfill(2)
        self.coupler.append(c_now)
        self.coupler=list(set(self.coupler))
        return c_now
    
    def readline_(self):#B:R
        straw=[]
        for i in self.qubits:
            straw.append('R'+str(1+i//6).zfill(2))
        self.readline=list(set(straw))
        return self.readline
    
    def barrier(self):#返回Barrier全部指令
        q_all=' '.join("Q{:02d}".format(num) for num in self.qubits)
        coupler_all= ' '.join([str(num) for num in self.coupler])
        self.readline_()
        readline_all=' '.join([str(num) for num in self.readline])
        if coupler_all=='':
            return f'B {q_all} {readline_all}'
        else:
            return f'B {q_all} {coupler_all} {readline_all}'
    
    def get_gate(self, op):#返回QCIS指令
        #单比特门
        if op.name == 'Identity':
            self.instruction.append(f'I Q{self.qubits[op.wires.tolist()[0]]:02d} 80')
        elif op.name == 'PauliX':
            self.instruction.append(f'X Q{self.qubits[op.wires.tolist()[0]]:02d} 3.1415926')
        elif op.name == 'PauliY':
            self.instruction.append(f'Y Q{self.qubits[op.wires.tolist()[0]]:02d} 3.1415926')
        elif op.name == 'PauliZ':
            self.instruction.append(f'Z Q{self.qubits[op.wires.tolist()[0]]:02d} 3.1415926')        
        elif op.name == "RX":
            self.instruction.append(f'RX Q{self.qubits[op.wires.tolist()[0]]:02d} {op.parameters[0]}')
        elif op.name == "RY":
            self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[0]]:02d} {op.parameters[0]}')
        elif op.name == "RZ":
            self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[0]]:02d} {op.parameters[0]}')       
        elif op.name == "Rot":
            self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[0]]:02d} {op.parameters[0]}')
            self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[0]]:02d} {op.parameters[1]}')
            self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[0]]:02d} {op.parameters[2]}')
        elif op.name == 'Hadamard':
            self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[0]]:02d} 3.1415926')
            self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[0]]:02d} 1.5707963')
        elif op.name == 'S':
            self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[0]]:02d} 1.5707963')
        elif op.name == 'T':
            self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[0]]:02d} 0.7853982')
            # print(self.instruction)
        #两比特门
        elif op.name == 'CZ':
            self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
            self.instruction.append(self.barrier())
            self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        elif op.name == 'CNOT':
            self.instruction.append(f'Y2M Q{self.qubits[op.wires.tolist()[1]]:02d}')
            self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
            self.instruction.append(self.barrier())
            self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
            self.instruction.append(f'Y2P Q{self.qubits[op.wires.tolist()[1]]:02d}')
        elif op.name == 'SWAP':
            self.instruction.append(f'Y2M Q{self.qubits[op.wires.tolist()[1]]:02d}')
            self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
            self.instruction.append(self.barrier())
            self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
            self.instruction.append(f'Y2P Q{self.qubits[op.wires.tolist()[1]]:02d}')

            self.instruction.append(f'Y2M Q{self.qubits[op.wires.tolist()[0]]:02d}')
            self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
            self.instruction.append(self.barrier())
            self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
            self.instruction.append(f'Y2P Q{self.qubits[op.wires.tolist()[0]]:02d}')

            self.instruction.append(f'Y2M Q{self.qubits[op.wires.tolist()[1]]:02d}')
            self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
            self.instruction.append(self.barrier())
            self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
            self.instruction.append(f'Y2P Q{self.qubits[op.wires.tolist()[1]]:02d}')
        # elif op.name == 'ISWAP':#
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[0]]:02d} 1.5707963')
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} 1.5707963')
        #     self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
        #     self.instruction.append(self.barrier())
        #     self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[0]]:02d} -1.5707963')
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} -1.5707963')
        #     self.instruction.append(self.barrier())
        #     self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        #     self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[0]]:02d} 1.5707963')
        #     self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[1]]:02d} 1.5707963')
        # elif op.name == 'CRX':
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} {op.parameters[0]/2}')

        #     self.instruction.append(f'Y2M Q{self.qubits[op.wires.tolist()[1]]:02d}')
        #     self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
        #     self.instruction.append(self.barrier())
        #     self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        #     self.instruction.append(f'Y2P Q{self.qubits[op.wires.tolist()[1]]:02d}')

        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} {-op.parameters[0]/2}')
            
        #     self.instruction.append(f'Y2M Q{self.qubits[op.wires.tolist()[1]]:02d}')
        #     self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
        #     self.instruction.append(self.barrier())
        #     self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        #     self.instruction.append(f'Y2P Q{self.qubits[op.wires.tolist()[1]]:02d}')
        # elif op.name == 'CRY':#不知道是X还是Y写错了。。。算一下
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} {op.parameters[0]/2}')
        #     self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[1]]:02d} -1.5707963')
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} 1.5707963')
        #     self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
        #     self.instruction.append(self.barrier())
        #     self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        #     self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[1]]:02d} 1.5707963')
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} -1.5707963')
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} {-op.parameters[0]/2}')
        #     self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[1]]:02d} -1.5707963')
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} 1.5707963')
        #     self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
        #     self.instruction.append(self.barrier())
        #     self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        #     self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[1]]:02d} 1.5707963')
        #     self.instruction.append(f'RY Q{self.qubits[op.wires.tolist()[1]]:02d} -1.5707963')
        # elif op.name == 'CRZ':
        #     self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[1]]:02d} {op.parameters[0]/2}')
        #     self.chip_topology(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])
        #     self.instruction.append(self.barrier())
        #     self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        #     self.instruction.append(f'RZ Q{self.qubits[op.wires.tolist()[1]]:02d} {-op.parameters[0]/2}')
        #     self.instruction.append(self.barrier())
        #     self.instruction.append(f'CZ {self.coupler_(self.qubits[op.wires.tolist()[0]], self.qubits[op.wires.tolist()[1]])}')
        else:
            raise ValueError(f"Gate {op.name} is not supported")
    def apply_observable(self, msmt):
    # 处理概率测量 (qml.probs)
        if isinstance(msmt, qml.measurements.ProbabilityMP):
            wires = msmt.wires.tolist()
            self.instruction.append(self.barrier())
            instruction = ['M']
            for wire in wires:
                instruction.append(f'Q{self.qubits[wire]:02d}')
            self.instruction.append(' '.join(instruction))
            return  # 概率测量无需旋转操作
    
    # 处理其他观测量（如期望值）
        obs = msmt.obs
        factors = []
        if hasattr(obs, "name") and obs.name == "Prod":
            factors = obs.operands
        else:
            factors = [obs]

        for factor in factors:
            if not hasattr(factor, "wires"):
                raise ValueError("Observable must act on specific wires.")
            wire = factor.wires.tolist()[0]
            if factor.name == 'PauliX':
                self.instruction.append(f'Y2M Q{self.qubits[wire]:02d}')
            elif factor.name == 'PauliY':
                self.instruction.append(f'X2P Q{self.qubits[wire]:02d}')
            elif factor.name == 'PauliZ':
                pass
            else:
                raise ValueError(f"Observable {factor.name} is not supported")

    # 生成联合测量指令
        wires = obs.wires.tolist()
        self.instruction.append(self.barrier())
        instruction = ['M']
        for wire in wires:
            instruction.append(f'Q{self.qubits[wire]:02d}')
        self.instruction.append(' '.join(instruction))    
    # def apply_observable(self, msmt):
    # # 获取观测量的所有因子（例如 PauliX(0) @ PauliY(1) 分解为 [PauliX(0), PauliY(1)]）
    #     obs = msmt.obs
    #     factors = []
    #     if hasattr(obs, "name") and obs.name == "Prod":  # 处理张量积
    #         factors = obs.operands
    #     else:
    #         factors = [obs]  # 单个观测量

    # # 对每个因子生成旋转指令
    #     for factor in factors:
    #         if not hasattr(factor, "wires"):
    #             raise ValueError("Observable must act on specific wires.")
    #         wire = factor.wires.tolist()[0]
    #         if factor.name == 'PauliX':
    #             self.instruction.append(f'Y2M Q{self.qubits[wire]:02d}')
    #         elif factor.name == 'PauliY':
    #             self.instruction.append(f'X2P Q{self.qubits[wire]:02d}')
    #         elif factor.name == 'PauliZ':
    #             pass  # Z测量无需旋转
    #         else:
    #             raise ValueError(f"Observable {factor.name} is not supported")

    # # 生成联合测量指令
    #     wires = obs.wires.tolist()  # 所有涉及的量子比特
    #     self.instruction.append(self.barrier())  # 添加屏障确保同步
    #     instruction = ['M']
    #     for wire in wires:
    #         instruction.append(f'Q{self.qubits[wire]:02d}')
    #     self.instruction.append(' '.join(instruction))

    
    def qcis(self,circuit: QuantumScript,
        execution_config = ExecutionConfig(
            grad_on_execution = False,
            use_device_gradient = True,
            use_device_jacobian_product = False,
            gradient_method = "adjoint",
            ), 
    ):#生成完整QCIS
        for i in range(len(self.qubits)):
            self.check_qubit(self.qubits[i])  # 检查qubit是否在范围内
        operations = circuit.operations
        for op in operations:
            self.get_gate(op)
        for msmt in circuit.measurements:#怎么把这个circuit传进来，还有那些内置属性咋整
            self.apply_observable(msmt)
        ins = '\n'.join(self.instruction)
        if self.printOP:
            print(ins)
        return ins

    # def run(self,msmt):
    #     if isinstance(msmt, qml.measurements.SampleMeasurement):
    #         self.datatype = "EVEVT_STATE"  # 设置数据类型为事件状态 
    #     qcis_copy=self.qcis().copy()
    #     circuit = qcis_copy
    #     qubits = self.qubits
    #     sampling_interval = 200e-6
    #     num_shots = self.shots
    #     d = RunCircuits(
    #         qubits=qubits,
    #         use_template=False,
    #         circuits=([circuit] * 100),
    #         data_type=self.datatype,
    #         sampling_interval=sampling_interval,
    #         num_shots=num_shots,
    #     )
    #     d.wait()
    #     frames = d.dataset.get_stream_data()
    #     return frames      
    
    
    # def calculate_average_variance(self,probs_list):
    #     if not probs_list:
    #         return 0.0
    
    #     n = int(math.log2(len(probs_list[0])))
    #     m = len(probs_list)
    
    #     # 初始化每个比特的期望值列表
    #     expectations = [[] for _ in range(n)]
    
    #     for probs in probs_list:
    #         for j in range(n):
    #             e_j = 0.0
    #             for k, p in enumerate(probs):
    #                 bit = (k >> j) & 1  # 获取第j比特的值
    #                 e_j += p * bit
    #             expectations[j].append(e_j)
    
    #     total_variance = 0.0
    #     for j in range(n):
    #         data = expectations[j]
    #         mean = sum(data) / m
    #         variance = sum((x - mean) ** 2 for x in data) / m
    #         total_variance += variance
    
    #     return (expectations,total_variance / n)

    #下面俩函数都是返回的pennylane上计算得到的值(对实验得到的数据后处理)
    # def apply_msmt(self, msmt):
    #     """处理期望值（expval）和方差（var）、概率分布(probs)测量"""
    #     str_list = self.run()
    #     m=self.calculate_average_variance(str_list)
    #     if isinstance(msmt, qml.measurements.ExpectationMP):
    #         return m[0]
    #     elif isinstance(msmt, qml.measurements.VarianceMP):
    #         return m[1] 
    #     elif isinstance(msmt, qml.measurements.ProbabilityMP):
    #         return str_list    
    #     elif isinstance(msmt, qml.measurements.SampleMeasurement):
    #         int_list_of_lists = [list(map(int, s)) for s in str_list]
    #         return int_list_of_lists
    #     else:
    #         raise ValueError(f"Measurement {msmt.__class__.__name__} not supported.")
    def execute(
        self,
        circuits: QuantumScriptOrBatch,
        execution_config = ExecutionConfig(
            grad_on_execution = False,
            use_device_gradient = True,
            use_device_jacobian_product = False,
            gradient_method = "adjoint",
            ), 
    ):
        # print(execution_config)
        if isinstance(circuits, qml.tape.QuantumScript):
            return self.qcis(circuits, execution_config)
        else:
            return tuple(self.qcis(c, execution_config) for c in circuits)
        
    # def compute_derivatives(self, 
    #     circuits, 
    #     execution_config = ExecutionConfig(
    #         grad_on_execution = False,
    #         use_device_gradient = True,
    #         use_device_jacobian_product = False,
    #         gradient_method = "adjoint",
    #         )
    # ):
    #     """Parameter Shift Rule"""
    #     gradients = []
    #     for circuit in circuits:
    #         grad = []
    #         for param_idx in circuit.trainable_params:
    #             # 创建偏移电路
    #             shifted_plus = self.create_shifted_circuit(circuit, param_idx, +np.pi/2)
    #             shifted_minus = self.create_shifted_circuit(circuit, param_idx, -np.pi/2)
                
    #             # 执行并计算梯度
    #             f_plus = self.execute(shifted_plus)
    #             f_minus = self.execute(shifted_minus)
    #             grad.append((f_plus - f_minus) / 2)
    #         gradients.append(np.array(grad))
    #     return tuple(gradients)

    # def create_shifted_circuit(self, circuit, param_idx, shift):
    #     """创建参数偏移后的电路"""
    #     params = circuit.get_parameters(trainable_only=False)
    #     params[param_idx] += shift
    #     return circuit.bind_new_parameters(params, circuit.trainable_params)
# A=Quanta_CS([0,65])
#m=A.chip_topology(0,65)
#f=A.readline()
#print(f)